import sys
from pylab import *
from pandas import DataFrame

def reflectance(nr=1.3, ni=0.1, q=pi/4.):
    """
    Calculate the reflected electric field from equation 20-61 (Fresnel
    Equations, University of Texas at Arlington.
    nr : real part of refractive index (float)
    ni : impaginary part of refractive index (float)
    q : angle of incidence (float)
    return : dictionary containing magnitude of reflected electric field 
    """ 
    dn2 = nr**2 - ni**2
    g1 = dn2 - (sin(q))**2

    TE = (cos(q) - sqrt(g1 + 2j*nr*ni))  /  \
         (cos(q) + sqrt(g1 + 2j*nr*ni))

    TM =  ( (dn2 + 2j*nr*ni) * cos(q) - sqrt(g1 + 2j*nr*ni) ) /    \
          ( (dn2 + 2j*nr*ni) * cos(q) + sqrt(g1 + 2j*nr*ni) ) 

    RE = (abs(TE))**2
    RM = (abs(TM))**2

    return  {'nr':nr, 'ni':ni, 'q':q, 'TE':TE, 'TM':TM, 'RE':RE, 'RM':RM}

def reflectance_w(um=1., e1m=1.1, e2m=0.04, ud=1, ed=1.1, q=pi/4.):
    """
    Calculate the relected electric field from 4.1-4.3, 4.8, 4.12-4.13, 4.18-4.25
    Whitaker J. Phys A: Math Gen. Vol 12, No.3 1979
    um : permeability of metal
    e1m : real dielectric metal (unitless)
    e2m : imaginary dielectric metal (unitless)
    ud : permeability of dielectric material
    ed : dielectric of dielectric material
    q : angle of incidence (float)
    return : dictionary containing magnitude of reflected electric field 
    """

    """ eq 3.3 """
    S = (e1m*um - ed*ud*(sin(q))**2)**2 + e2m**2*um**2

    """ s polarization (equations 4.1 - 4.3) """
    u_p = (S**0.5 + (e1m*um - ed*md*(sin(q))**2))**0.5
    u_m = (S**0.5 - (e1m*um - ed*md*(sin(q))**2))**0.5
    t1 = um**2*ed *(cos(q))**2 + ud*S**0.5 + 2**0.5*um*ud**0.5*ed**0.5*u_p*cos(q)
    Erefl_s = (um**2*ed*(cos(q))**2 - ud*S**0.5 - (2j)**0.5*um*ud**0.5*ed**0.5*u_m*cos(q)) / t1

    """ equation 4.8 """
    Erefr_s = 2*um*ed**0.5*cos(q)*((um*ed**0.5*cos(q) + 2**(-0.5)*ud**0.5*u_p) - (2j)**(-0.5)*ud**0.5*u_m) / t1


    """ p polarization (equation 4.12 - 4.13) """
    t3 = ud*um*(e1m**2 + e2m**2)*(cos(q))**2 + um*ed*S**0.5 + 2**0.5*ud**0.5*ed**0.5*u_p*cos(q)*(S**0.5 + ed*ud*(sin(q)))**2
    Erefl_p = ((ud*um*(e1m**2 + e2m**2)*(cos(q))**2 - um*ed*S**0.5) + \
            1j*(2**0.5*ud**0.5*ed**0.5*cos(q)*u_m*(S**0.5 - ed*ud*(sin(q))**2))) / t3

    """ equation 4.18 """
    Erefr_px = (2*um*ed*S**0.5 + 2**0.5*ud**0.5*ed**0.5*cos(q)*(u_p*(S**0.5 + ed*ud*(sin(q))**2) + \
            1j*u_m*(ed*ud*(sin(q))**2 - S**0.5))) / t3

    """ equation 4.22 """
    Erefr_pz = 2*ud**0.5*um*ed*cos(q)*sin(q)*((ud**0.5*e1m*cos(q) + ed**0.5*2**(-0.5)*u_p) - \
            1j*(ud**0.5*e2m*cos(q) + ed**0.5*2**(-0.5)*u_m)) / t3

    return {'um':um, 'e1m':e1m, 'e2m':e2m, 'ud':ud, 'q':q, 
            'Erefl_s':Erefl_s, 'Erefr_s':Erefr_s, 
            'Erefl_p':Erefl_p, 'Erefr_px':Erefr_px, 'Erefr_pz':Erefr_pz}

def plot(nr=1.3, ni=0.1):
    """
    Plot the reflectance.
    nr : real part of refractive index (float)
    ni : impaginary part of refractive index (float)
    """
    angle = linspace(0.01, pi/2, 100)
    refl = DataFrame([reflectance(nr, ni, q) for q in angle])
    fig = figure()
    ax = [refl.plot(x='q',y=y) for y in ['RE','RM']]
    ax[0].set_xlabel(r'$\theta$ (rad)')
    ax[0].set_ylabel(r'R')
    ax[1].legend(['RE','RM'])
    show()
    return ax
    
if __name__=='__main__':
    nr = float(sys.argv[1])
    ni = float(sys.argv[2])
    ioff()
    ax = plot(nr=nr, ni=ni)
